from AStock.ASQuery import ASQuery_nearest_trade_date, ASQuery_stock_day


class ASEarningRatio(object):
    def __init__(self, codes, begin_date, end_date):
        self._codes = codes
        self._begin_date = begin_date
        self._end_date = end_date

    def run(self, save_to_db=True, save_to_file=False, sort=False):
        codes = self._codes
        begin_date = self._begin_date
        end_date = self._end_date
        # TODO 权重暂未考虑
        # weights = args.weights

        # if weights and len(weights) != len(codes):
        #     print('codes and weights not match')
        #   return

        begin_date = ASQuery_nearest_trade_date(begin_date)
        end_date = ASQuery_nearest_trade_date(end_date)
        df_begin_price = ASQuery_stock_day(codes, begin_date, begin_date, fields=['code', 'date', 'close'], type='qfq')
        df_end_price = ASQuery_stock_day(codes, end_date, end_date, fields=['code', 'date', 'close'], type='qfq')
        df = df_begin_price.join(df_end_price.set_index('code'), on='code', how='left', lsuffix='_begin',
                                 rsuffix='_end')
        # df columns: code date_begin date_end close_begin close_end
        df['earning_ratio'] = (df['close_end'] - df['close_begin']) / df['close_begin']
        if sort:
            df = df.sort_values(by='earning_ratio')
        else:
            # 按 codes 输入顺序排序
            codes_sort_dict = {code: i for i, code in enumerate(codes)}
            df['codes_sort'] = df.apply(lambda row: codes_sort_dict[row['code']], axis=1)
            df = df.sort_values(by='codes_sort')
            df = df.drop('codes_sort', axis=1)
        # df.to_excel('D://er.xlsx')
        print(df)
        print('mean earning ratio: {}%'.format(round(df['earning_ratio'].mean() * 100, 2)))
